| 31 | |
| 32 | |
| 33 | class QueryForm(ModelForm): |
| 34 | |
| 35 | sql = SqlField() |
| 36 | snapshot = BooleanField(widget=CheckboxInput, required=False) |
| 37 | few_shot = BooleanField(widget=CheckboxInput, required=False) |
| 38 | database_connection = CharField(widget=Select, required=False) |
| 39 | |
| 40 | def __init__(self, *args, **kwargs): |
| 41 | super().__init__(*args, **kwargs) |
| 42 | self.fields["database_connection"].widget.choices = self.connections |
| 43 | if not self.instance.database_connection: |
| 44 | default_db = default_db_connection() |
| 45 | self.initial["database_connection"] = default_db_connection().alias if default_db else None |
| 46 | self.fields["database_connection"].widget.attrs["class"] = "form-select" |
| 47 | |
| 48 | def clean(self): |
| 49 | # Don't overwrite created_by_user |
| 50 | if self.instance and self.instance.created_by_user: |
| 51 | self.cleaned_data["created_by_user"] = \ |
| 52 | self.instance.created_by_user |
| 53 | return super().clean() |
| 54 | |
| 55 | def clean_database_connection(self): |
| 56 | connection_id = self.cleaned_data.get("database_connection") |
| 57 | if connection_id: |
| 58 | try: |
| 59 | return DatabaseConnection.objects.get(id=connection_id) |
| 60 | except DatabaseConnection.DoesNotExist as e: |
| 61 | raise ValidationError("Invalid database connection selected.") from e |
| 62 | return None |
| 63 | |
| 64 | @property |
| 65 | def created_at_time(self): |
| 66 | return self.instance.created_at.strftime("%Y-%m-%d") |
| 67 | |
| 68 | @property |
| 69 | def connections(self): |
| 70 | default_db = default_db_connection() |
| 71 | if default_db is None: |
| 72 | return [] |
| 73 | |
| 74 | # Ensure the default connection appears first in the dropdown in the form |
| 75 | result = DatabaseConnection.objects.annotate( |
| 76 | custom_order=Case( |
| 77 | When(id=default_db_connection().id, then=Value(0)), |
| 78 | default=Value(1), |
| 79 | output_field=IntegerField(), |
| 80 | ) |
| 81 | ).order_by("custom_order", "id") |
| 82 | return [(c.id, c.alias) for c in result.all()] |
| 83 | |
| 84 | class Meta: |
| 85 | model = Query |
| 86 | fields = ["title", "sql", "description", "snapshot", "database_connection", "few_shot"] |