Made new urls not generate a new token if they already exist in the database

tiny0/models.py

@@ -3,7 +3,7 @@ from tiny0 import db
 class URL(db.Model):
 	id = db.Column(db.Integer, primary_key=True)
 	token = db.Column(db.String(6), index=True, unique=True, nullable=False)
-	url = db.Column(db.String(2000), nullable=False)
+	url = db.Column(db.String(2000), index=True, unique=True, nullable=False)
 
 	def __repr__(self):
 		return f"'{self.id}' '{self.token}' '{self.url}'"

tiny0/routes.py

@@ -13,6 +13,9 @@ def index():
 
 	# If the form was valid
 	if form.validate_on_submit():
+
+		# If the given url is a rick roll
+		if ("youtube.com/watch?v=dQw4w9WgXcQ" in form.url.data) or ("youtu.be/dQw4w9WgXcQ" in form.url.data):
 			# Generate a valid token
 			token = gen_valid_token()
 
@@ -23,7 +26,27 @@ def index():
 			# Return the url page with the shortened url
 			return render_template("url.html", url=WEBSITE_DOMAIN + "/" + token)
 
-	# If the form was invalid or not submitted
+		# Query the urls that are not a rick roll in the database
+		query = URL.query.filter_by(url=form.url.data).first()
+
+		# If the url exists in the database
+		if query:
+			# Return the url page with the previously shortened url
+			return render_template("url.html", url=WEBSITE_DOMAIN + "/" + query.token)
+
+		# Else if the url does not exist in the database
+		else:
+			# Generate a valid token
+			token = gen_valid_token()
+
+			# Add the token and the given url to the database
+			db.session.add(URL(token=token, url=form.url.data))
+			db.session.commit()
+
+			# Return the url page with the shortened url
+			return render_template("url.html", url=WEBSITE_DOMAIN + "/" + token)
+
+	# Else if the form was invalid or not submitted
 	else:
 		# Return the index page with the form
 		return render_template("index.html", form=form)

tiny0/token.py

@@ -1,3 +1,4 @@
+from tiny0 import db
 from tiny0.models import URL
 from secrets import choice
 
@@ -10,7 +11,7 @@ def gen_valid_token():
 		token = "".join(choice(token_characters) for i in range(6))
 
 		# If the token does not exists in the database
-		if not URL.query.filter_by(token=token).first():
+		if not db.session.query(db.session.query(URL).filter_by(token=token).exists()).scalar():
 			# Break the loop
 			break